import requests
import base64
import os

# --- 配置区域 ---

# 1. 您的腾讯云函数服务器 URL (请确保路径正确)
# 这里的路径需要是完整的 API 路径
SERVER_URL = "https://1252410844-i0l6vn4s0r.ap-tokyo.tencentscf.com/api/generate-image-from-image"

# 2. 定义输入和输出文件路径
INPUT_IMAGE_PATH = "fbb.jpg"
OUTPUT_IMAGE_PATH = "generated_from_server.png" # 服务器返回的图片将保存为此文件名

# 3. [可选] 您可以自定义每次测试的提示词
# 如果不提供，服务器会使用代码中设置的默认提示词
CUSTOM_PROMPT = "create a 1/7 scale commercialized figurine of the characters in the picture, in a realistic style, in a real environment. The figurine is placed on a computer desk. The figurine has a round transparent acrylic base, with no text on the base. The content on the computer screen is the Zbrush modeling process of this figurine. Next to the computer screen is a BANDAI-style toy packaging box printed with the original artwork. The packaging features two-dimensional flat illustrations."


# --- 脚本执行区域 ---

def run_test():
    """
    执行图片上传、生成和下载的完整测试流程。
    """
    print(f"准备测试服务器: {SERVER_URL}")

    # 步骤 1: 检查本地图片文件是否存在
    if not os.path.exists(INPUT_IMAGE_PATH):
        print(f"❌ 错误：找不到测试图片 '{INPUT_IMAGE_PATH}'。请确保它和脚本在同一个目录下。")
        return

    print(f"✅ 找到了本地图片: {INPUT_IMAGE_PATH}")

    # 步骤 2: 准备上传的数据
    # 'multipart/form-data' 是文件上传时常用的格式
    # 'files' 字典的 key 'userImage' 必须和您 server.js 中 upload.single('userImage') 的参数完全一致
    try:
        with open(INPUT_IMAGE_PATH, 'rb') as image_file:
            files_to_upload = {
                'userImage': (os.path.basename(INPUT_IMAGE_PATH), image_file, 'image/jpeg')
            }
            # 'data' 字典用于传递额外的表单字段，比如我们的提示词
            form_data = {
                'prompt': CUSTOM_PROMPT
            }

            print("🚀 正在上传图片并请求 AI 生成新图片，这可能需要几十秒，请耐心等待...")

            # 步骤 3: 发送 POST 请求
            # 注意：腾讯云函数可能有冷启动延迟，所以设置一个较长的超时时间（例如60秒）
            response = requests.post(SERVER_URL, files=files_to_upload, data=form_data, timeout=60)

            # 步骤 4: 处理服务器的响应
            # 检查 HTTP 状态码是否为 200 (OK)
            response.raise_for_status() 

            print("✅ 服务器成功响应！正在处理返回的数据...")

            # 解析返回的 JSON 数据
            result = response.json()

            # 检查返回的数据格式是否正确
            if result.get('success') and result.get('imageData'):
                base64_string = result['imageData'].get('base64')
                
                if not base64_string:
                    print("❌ 错误：服务器返回成功，但没有找到 base64 图片数据。")
                    return

                # 解码 Base64 字符串变回二进制图片数据
                generated_image_data = base64.b64decode(base64_string)

                # 将图片数据写入新文件
                with open(OUTPUT_IMAGE_PATH, 'wb') as output_file:
                    output_file.write(generated_image_data)

                print(f"\n🎉🎉🎉 测试成功！新图片已保存为: {OUTPUT_IMAGE_PATH}")
            else:
                print(f"❌ 错误：服务器返回的数据格式不正确。服务器消息: {result.get('error', '无详细错误信息')}")

    except requests.exceptions.Timeout:
        print("❌ 请求超时：服务器在60秒内没有响应。可能是冷启动时间过长或网络问题，请稍后重试。")
    except requests.exceptions.RequestException as e:
        print(f"❌ 网络请求失败：{e}")
        # 如果有响应内容，也打印出来方便排查
        if e.response is not None:
            print(f"服务器状态码: {e.response.status_code}")
            print(f"服务器原始响应: {e.response.text}")
    except Exception as e:
        print(f"❌ 发生未知错误: {e}")


if __name__ == "__main__":
    run_test()